import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.models import *
from tensorflow.keras.layers import *
import regex as re
import scipy as sp
from tensorflow.keras.utils import to_categorical
from multiprocessing import Process

# import tensorflow as tf
# config = tf.compat.v1.ConfigProto()
# config.gpu_options.allow_growth = True  
# sess = tf.compat.v1.Session(config=config)

def seq2oh(seq):
    vocab = {
        'A': [1,0,0,0,0],
        'T': [0,1,0,0,0],
        'C': [0,0,1,0,0],
        'G': [0,0,0,1,0],
        '_': [0,0,0,0,1]
    }
    if len(seq)<100:
        seq = '_' * (100-len(seq)) + seq
    oh = []
    for item in seq:
        oh.append(vocab[item])
    return np.array(oh)


def oh2seq(oh):
    seq = ''
    vocab = {
        'A': [1,0,0,0,0],
        'T': [0,1,0,0,0],
        'C': [0,0,1,0,0],
        'G': [0,0,0,1,0],
        '_': [0,0,0,0,1]
    }
    for item in oh:
        char = list(vocab.keys())[list(vocab.values()).index(list(item))]
        seq += char
    seq = seq.split('_')[-1]
    return seq


def preprocessing_seq(fasta_text):
    name_lst = []
    seq_lst = []
    str_lst = fasta_text.splitlines(keepends=False)
    for i in range(len(str_lst)-1):
        if str_lst[i][0] == '>':
            seq = str_lst[i+1].upper()
            seq = seq.replace('U','T')
            if len(seq) > 100:
                seq = seq[-100:]
            if len(seq) >= 25:
                seq_lst.append(seq)
                name_lst.append(str_lst[i][1:])
    return seq_lst, name_lst


def predict_rl(text_fa, is_gen, task_num):
    import os 
    os.environ['CUDA_VISIBLE_DEVICES']='2'
    gpu_list = tf.config.experimental.list_physical_devices('GPU')
    if len(gpu_list) > 0 :
        for gpu in gpu_list:
            try:
                tf.config.experimental.set_memory_growth(gpu, True)
                print(gpu)
            except RuntimeError as e:
                print(e)
    else:
    	print("Got no GPUs")
    seq_lst, name_lst = preprocessing_seq(text_fa)
    seq_oharr = []
    for item in seq_lst:
        seq_f = seq2oh(item)
        seq_oharr.append(seq_f)
    seq_oharr = np.array(seq_oharr)
    eval_model_25 = load_model('utrgen_app/utrgen_utils/models/CGRU_25_100.hd5')
    if is_gen:
        eval_model_50 = load_model('utrgen_app/utrgen_utils/models/CGRU_50.hd5')
        gen_model = load_model('utrgen_app/utrgen_utils/models/UTR_genAE_CGRU_25_100_compiled.hd5')
        opt_lst = []
        opt_oh_lst = []
        pred_mrl_ori = eval_model_25.predict(seq_oharr, verbose=0).flatten()
        seq_ae = gen_model.predict(seq_oharr, verbose=0)
        for i in range(len(seq_lst)):
            ae_arg = np.argmax(seq_ae[i], axis=-1)
            ae_argoh = to_categorical(ae_arg, 5)
            for m in range(100-len(seq_lst[i]),100):
                if ae_argoh[m][4] == 1:
                    ae_argoh[m] = np.array([0,0,1,0,0])
            seq_nat = oh2seq(ae_argoh)
            seq_nat_opt = ''
            for j in range(len(seq_lst[i])):
                try:
                    if seq_lst[i][j] != seq_nat[j]:
                        seq_nat_opt = seq_nat_opt + seq_nat[j].lower()
                    else:
                        seq_nat_opt = seq_nat_opt + seq_nat[j]
                except:
                    pass
            opt_oh_lst.append(ae_argoh)
            opt_lst.append(seq_nat_opt)
        pred_mrl_opt = eval_model_25.predict(np.array(opt_oh_lst), verbose=0).flatten()
        feat_arr_50 = np.array(opt_oh_lst)
        feat_arr_50[:,:50] = np.array([0, 0, 0, 0, 1])
        pred_mrl_opt_50 = eval_model_50.predict(np.array(feat_arr_50), verbose=0).flatten()
        inc_mrl = pred_mrl_opt - pred_mrl_ori
        df_result = pd.DataFrame({'Sequence name': name_lst,
                                  'Original UTR': seq_lst, 
                                  'Original RL': pred_mrl_ori, 
                                  'Optimized UTR': opt_lst,
                                  'Optimized RL': pred_mrl_opt, 
                                  'Optimized RL_50nt': pred_mrl_opt_50,
                                  'Increased RL': inc_mrl, 
                                 })
    else:
        pred_mrl_ori = eval_model_25.predict(seq_oharr, verbose=0).flatten()
        df_result = pd.DataFrame({'Sequence name': name_lst,
                                  'Original UTR': seq_lst, 
                                  'Predicted RL': pred_mrl_ori, 
                                 })
    file_name = 'utrgen_app/static/output/{}_report.xlsx'.format(task_num)
    df_result.to_excel(file_name, index=False)
    print('Analysis completed! You can find output file in {}.'.format(file_name))
    return df_result


# def output_result(seq_fa, is_gen=False, task_num=None):
#     pool = Pool(processes=2)
#     res_df = pool.apply(predict_rl,(seq_fa, is_gen))
#     pool.close()
#     pool.join()
#     # t = str(int(time.time()))
#     file_name = 'utrgen_app/static/output/{}_report.xlsx'.format(task_num)
#     res_df.to_excel(file_name, index=False)
#     print('Analysis completed! You can find output file in {}.'.format(file_name))
#     return 0





